package pers.whj.util.tree;

public class RedBlackTreeImpl<T extends Comparable<T>> extends BinarySearchTreeImpl<T> implements RedBlackTree<T> {
    public static final boolean BLACK = true;
    public static final boolean RED = false;

    public RedBlackTreeImpl() {
    }

    @Override
    public T precursor(T data) {
        Node<T> node = (Node<T>) search(this.root, data);
        checkNode(node);
        node = precursor(node);
        return node == null ? null : node.data;
    }

    @Override
    public T successor(T data) {
        Node<T> node = (Node<T>) search(this.root, data);
        checkNode(node);
        node = successor(node);
        return node == null ? null : node.data;
    }

    @Override
    public boolean insert(T data) {
        checkInsertData(data);
        return this.insert(new Node<>(data, BLACK));
    }

    @Override
    public boolean remove(T data) {
        Node<T> node = (Node<T>) search(this.root, data);
        return node != null && this.remove(node);
    }

    protected static class Node<T extends Comparable<T>> extends AbstractTree.Node<T> {
        protected boolean color;
        protected Node<T> parent;

        public Node(T data, boolean color) {
            super(data);
            this.color = color;
        }

        public Node(T data, boolean color, Node<T> left, Node<T> right, Node<T> parent) {
            super(data, left, right);
            this.color = color;
            this.parent = parent;
        }

        @Override
        public void reset() {
            super.reset();
            this.color = false;
            this.parent = null;
        }
    }

    protected boolean colorOf(Node<T> node) {
        return node == null ? BLACK : node.color;
    }

    protected boolean isRed(Node<T> node) {
        return node != null && node.color == RED;
    }

    protected boolean isBlack(Node<T> node) {
        return node == null || node.color == BLACK;
    }

    protected void setBlack(Node<T> node) {
        if (node != null) {
            node.color = BLACK;
        }
    }

    protected void setRed(Node<T> node) {
        if (node != null) {
            node.color = RED;
        }
    }

    protected void setParent(Node<T> node, Node<T> parent) {
        if (node != null) {
            node.parent = parent;
        }
    }

    protected void setColor(Node<T> node, boolean color) {
        if (node != null) {
            node.color = color;
        }
    }

    @Override
    protected AbstractTree.Node<T> parentOf(AbstractTree.Node<T> node) {
        return node == null ? null : ((Node<T>) node).parent;
    }

    @Override
    protected AbstractTree.Node<T> parentOf(AbstractTree.Node<T> parent, AbstractTree.Node<T> node) {
        return node == null ? null : ((Node<T>) node).parent;
    }

    protected Node<T> precursor(Node<T> node) {
        if (node == null) {
            return null;
        }
        if (node.left != null) {
            return (Node<T>) max(node.left);
        }
        Node<T> parent = node.parent;
        while (parent != null && node == parent.left) {
            node = parent;
            parent = parent.parent;
        }
        return parent;
    }

    protected Node<T> successor(Node<T> node) {
        if (node == null) {
            return null;
        }
        if (node.right != null) {
            return (Node<T>) min(node.right);
        }
        Node<T> parent = node.parent;
        while (parent != null && node == parent.right) {
            node = parent;
            parent = parent.parent;
        }
        return parent;
    }

    /**
     * insert a node into this tree
     *
     * @param node target node to insert
     */
    protected boolean insert(Node<T> node) {
        Node<T> x = (Node<T>) this.root;
        Node<T> y = null;
        int compare;
        while (x != null) {
            y = x;
            compare = node.data.compareTo(x.data);
            if (compare < 0) {
                x = (Node<T>) x.left;
            } else if (compare > 0) {
                x = (Node<T>) x.right;
            } else {
                return false;
            }
        }
        node.parent = y;
        if (y != null) {
            compare = node.data.compareTo(y.data);
            if (compare < 0) {
                y.left = node;
            } else {
                y.right = node;
            }
        } else {
            this.root = node;
        }
        node.color = RED;
        insertFixUp(node);
        this.size++;
        return true;
    }

    /**
     * remove the given node, this node should exist
     *
     * @param node given node
     * @return if success
     */
    protected boolean remove(Node<T> node) {
        Node<T> child, parent;
        boolean color;
        // if node has two children
        if (node.left != null && node.right != null) {
            // find the successor node to replace the remove node
            Node<T> replace = (Node<T>) node.right;
            while (replace.left != null) {
                replace = (Node<T>) replace.left;
            }
            // if node is root node
            Node<T> tmp = (Node<T>) parentOf(node);
            if (tmp == null) {
                this.root = replace;
            } else {
                if (tmp.left == node) {
                    tmp.left = replace;
                } else {
                    tmp.right = replace;
                }
            }
            // store the relevant data of replace node
            child = (Node<T>) replace.right;
            parent = (Node<T>) parentOf(replace);
            color = colorOf(replace);
            // if the remove node is the parent node of its successor node
            if (parent == node) {
                parent = replace;
            } else {
                if (child != null) {
                    setParent(child, parent);
                }
                parent.left = child;
                replace.right = node.right;
                setParent((Node<T>) node.right, replace);
            }
            replace.parent = node.parent;
            replace.color = node.color;
            replace.left = node.left;
            ((Node<T>) node.left).parent = replace;
            if (color == BLACK) {
                removeFixUp(child, parent);
            }
            node.reset();
            this.size--;
            return true;
        }
        if (node.left != null) {
            child = (Node<T>) node.left;
        } else {
            child = (Node<T>) node.right;
        }
        parent = node.parent;
        color = node.color;
        if (child != null) {
            child.parent = parent;
        }
        // if node is root node
        if (parent == null) {
            this.root = child;
        } else {
            if (parent.left == node) {
                parent.left = child;
            } else {
                parent.right = child;
            }
        }
        if (color == BLACK) {
            removeFixUp(child, parent);
        }
        node.reset();
        this.size--;
        return true;
    }

    @Override
    protected boolean remove(T data, AbstractTree.Node<T> node, AbstractTree.Node<T> parent) {
        return this.remove(data);
    }

    protected void rotateLeft(Node<T> node) {
        // assert node has right child
        Node<T> rightNode = (Node<T>) node.right;
        node.right = rightNode.left;
        if (rightNode.left != null) {
            ((Node<T>) rightNode.left).parent = node;
        }
        rightNode.parent = node.parent;
        if (node.parent == null) {
            this.root = rightNode;
        } else {
            if (node.parent.left == node) {
                node.parent.left = rightNode;
            } else {
                node.parent.right = rightNode;
            }
        }
        rightNode.left = node;
        node.parent = rightNode;
    }

    protected void rotateRight(Node<T> node) {
        // assert node has left child
        Node<T> leftNode = (Node<T>) node.left;
        node.left = leftNode.right;
        if (leftNode.right != null) {
            ((Node<T>) leftNode.right).parent = node;
        }
        leftNode.parent = node.parent;
        if (node.parent == null) {
            this.root = leftNode;
        } else {
            if (node.parent.left == node) {
                node.parent.left = leftNode;
            } else {
                node.parent.right = leftNode;
            }
        }
        leftNode.right = node;
        node.parent = leftNode;
    }

    /**
     * insertion correction method
     * after insertion, call this method to reshape it into a red-black tree
     *
     * @param node insertion node
     */
    protected void insertFixUp(Node<T> node) {
        Node<T> parent, grandParent;
        while ((parent = (Node<T>) parentOf(node)) != null && isRed(parent)) {
            grandParent = (Node<T>) parentOf(parent);
            if (parent == grandParent.left) {
                Node<T> uncle = (Node<T>) grandParent.right;
                // uncle node is red color
                if (uncle != null && isRed(uncle)) {
                    setBlack(uncle);
                    setBlack(parent);
                    setRed(grandParent);
                    node = grandParent;
                    continue;
                }
                // uncle node is black, current node is right child
                if (parent.right == node) {
                    rotateLeft(parent);
                    Node<T> tmp = parent;
                    parent = node;
                    node = tmp;
                }
                setBlack(parent);
                setRed(grandParent);
                rotateRight(grandParent);
            } else {
                Node<T> uncle = (Node<T>) grandParent.left;
                // uncle node is red color
                if (uncle != null && isRed(uncle)) {
                    setBlack(uncle);
                    setBlack(parent);
                    setRed(grandParent);
                    node = grandParent;
                    continue;
                }
                // uncle node is black, current node is left child
                if (parent.left == node) {
                    rotateRight(parent);
                    Node<T> tmp = parent;
                    parent = node;
                    node = tmp;
                }
                setBlack(parent);
                setRed(grandParent);
                rotateLeft(grandParent);
            }
        }
        setBlack((Node<T>) this.root);
    }

    /**
     * remove correction method
     * after remove a node, call this method to reshape it into a red-black tree
     *
     * @param node   replace node
     * @param parent the parent node of replace node
     */
    protected void removeFixUp(Node<T> node, Node<T> parent) {
        Node<T> other;
        while ((node == null || node.color == BLACK) && node != this.root) {
            if (parent.left == node) {
                other = (Node<T>) parent.right;
                // if its sibling is red
                if (isRed(other)) {
                    setBlack(other);
                    setRed(parent);
                    rotateLeft(parent);
                    other = (Node<T>) parent.right;
                }
                if ((other.left == null || ((Node<T>) other.left).color == BLACK)
                        && (other.right == null || ((Node<T>) other.right).color == BLACK)) {
                    setRed(other);
                    node = parent;
                    parent = (Node<T>) parentOf(node);
                } else {
                    if (other.right == null || ((Node<T>) other.right).color == BLACK) {
                        setBlack((Node<T>) other.left);
                        setRed(other);
                        rotateRight(other);
                        other = (Node<T>) parent.right;
                    }
                    setColor(other, colorOf(parent));
                    setBlack(parent);
                    setBlack((Node<T>) other.right);
                    rotateLeft(parent);
                    node = (Node<T>) this.root;
                    break;
                }
            } else {
                other = (Node<T>) parent.left;
                if (isRed(other)) {
                    setBlack(other);
                    setRed(parent);
                    rotateRight(parent);
                    other = (Node<T>) parent.left;
                }
                if ((other.left == null || ((Node<T>) other.left).color == BLACK)
                        && (other.right == null || ((Node<T>) other.right).color == BLACK)) {
                    setRed(other);
                    node = parent;
                    parent = (Node<T>) parentOf(node);
                } else {
                    if (other.left == null || ((Node<T>) other.left).color == BLACK) {
                        setBlack((Node<T>) other.right);
                        setRed(other);
                        rotateLeft(other);
                        other = (Node<T>) parent.left;
                    }
                    setColor(other, colorOf(parent));
                    setBlack(parent);
                    setBlack((Node<T>) other.left);
                    rotateRight(parent);
                    node = (Node<T>) this.root;
                    break;
                }
            }
        }
        if (node != null) {
            node.color = BLACK;
        }
    }
}